package com.fly.base.util;

import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.function.Predicate;

/**
 * tree和list互转工具类
 * todo 1.list转tree时加入一个回调参数，使得外部可以操控
 * todo 2.tree转list
 */
public class TreeUtils {

    /**
     * list转tree
     * 代码最简单的版本，使用了两次循环
     * @param source 原始列表
     * @param idFn 获取id的方法
     * @param pidFn 获取父id的方法
     * @param isRoot 判断是否是根节点
     * @param setChildren 设置子节点
     * @param <T> 元素类型
     * @param <R> id类型
     * @return tree
     */
    public static <T, R> List<T> list2Tree(List<T> source, Function<T, R> idFn, Function<T, R> pidFn,
                                           Predicate<T> isRoot, BiConsumer<T, List<T>> setChildren) {
        // 1.找出根节点，并将子节点根据父id分组
        List<T> rootList = new ArrayList<>();
        Map<R, List<T>> nodeMap = new HashMap<>();
        for (T t : source) {
            if (isRoot.test(t)) {
                rootList.add(t);
            } else {
                List<T> nodeList = nodeMap.getOrDefault(pidFn.apply(t), new ArrayList<>());
                nodeList.add(t);
                nodeMap.put(pidFn.apply(t), nodeList);
            }
        }
        // 遍历根节点，递归设置子节点
        rootList.forEach(e -> buildChildren(e, nodeMap, idFn, setChildren));
        return rootList;
    }

    /**
     * 递归查找子节点
     * @param current 当前节点
     * @param nodeMap 节点集合
     * @param idFn 获取id的方法
     * @param setChildren 设置子节点
     * @param <T> 元素类型
     * @param <R> id类型
     */
    private static <T, R> void buildChildren(T current, Map<R, List<T>> nodeMap,
                                             Function<T, R> idFn, BiConsumer<T, List<T>> setChildren) {
        List<T> children = nodeMap.get(idFn.apply(current));
        setChildren.accept(current, children);
        // 递归
        if (!CollectionUtils.isEmpty(children)) {
            children.forEach(e -> buildChildren(e, nodeMap, idFn, setChildren));
        }
    }
}
